(* Preprocess a Boolean program, return an equivalent program.
      The constructs that are changed while preprocessing is:
        - Global and local variables will be disjoint.
        - The return value of a procedure is never stored in an global variable.
 *)

open Ast;;

let tmpNum = ref 0;;

(* Sets the counter to default value (0)  *)
let resetTmpCounter () = tmpNum := 0; ();;

(* Create a new variable name, we have to make sure the new variable name is not a already declared variable. glob is the global variables, and all locals are the locally declared variables and parameters in the current procedure. *)
let rec createNewVar glob allLocals =
  tmpNum := (!tmpNum + 1);
  let varName = "tmp" ^ (string_of_int !tmpNum) in
    if ((not (Translation.varIsIn varName glob)) && (not (Translation.varIsIn varName allLocals))) then begin
      varName
    end
    else begin
      createNewVar glob allLocals
    end
  ;;

(* preprocesses a local declaration block. *)
let rec prepDecl variableSubs glob allLocals decl = 
  match decl with
    Declaration(labels, id) -> 
      if (Translation.varIsIn id glob) then (* the declared variable is a global var, make a substitution with a new variable name *)
        let newVariable = createNewVar glob allLocals in
          ((id, newVariable) :: variableSubs, Declaration(labels, newVariable))
      else 
        (variableSubs, Declaration(labels, id))
    | Declcomp(labels, d1, d2) -> 
        let (variableSubs1, newd1) = prepDecl variableSubs glob allLocals d1 in
        let (variableSubs2, newd2) = prepDecl variableSubs glob allLocals d2 in
          (variableSubs1 @ variableSubs2, Declcomp(labels, newd1, newd2))
    | Empty -> (variableSubs, decl)
	;;

(* Return the variable associated with var in variableSubs if possible, else returns var *)
let substituteVar var variableSubs =
  try List.assoc var variableSubs with
    Not_found -> var;;

(* Same as above but for a list of variables *)
let substituteVars varList variableSubs =
  let rec iter = function
    var :: rest -> (substituteVar var variableSubs) :: iter rest
    | [] -> []
  in
    iter varList;;
    
(* Substitute variables in an expression *)   
let rec substituteVarInExpr expr variableSubs =
  match expr with
  	Value(True) -> Value(True)
	| Value(False) ->  Value(False)
	| Variable(id) ->  Variable(substituteVar id variableSubs)
	| Not(e) ->  Not(substituteVarInExpr e variableSubs)
	| Equal(e1, e2) -> Equal(substituteVarInExpr e1 variableSubs, substituteVarInExpr e2 variableSubs)
	| And(e1, e2) -> And(substituteVarInExpr e1 variableSubs, substituteVarInExpr e2 variableSubs)
	| Or(e1, e2) -> Or(substituteVarInExpr e1 variableSubs, substituteVarInExpr e2 variableSubs)
	| Group(e) -> Group(substituteVarInExpr e variableSubs)
	;;

(* Same as above but for a list of expressions *)
let substituteVarInExprs exprs variableSubs =
  let rec iter = function
    e :: rest -> (substituteVarInExpr e variableSubs) :: iter rest
    | [] -> []
  in
    iter exprs;;

(* Substitute variables in the decider if the decider is an expression *)    
let substituteDecider d variableSubs =
  match d with
    Expression expr -> Expression(substituteVarInExpr expr variableSubs)
    | _ -> Undet
  ;;

(* Preprocesses a statement, and returns a new, equivalent, statement. *)
let rec prepStm variableSubs tmp glob s = 
  match s with
    | Skip(labels) -> Skip(labels)
    | Comp(labels, s1, s2) -> Comp(labels, prepStm variableSubs tmp glob s1, prepStm variableSubs tmp glob s2)
    | Callvoid(labels, id1, eList) -> Callvoid(labels, id1, substituteVarInExprs eList variableSubs)
    | Callbool(labels, id1, id2, eList) -> (* we are not allowed to assign a global variable the result of the procedure *)
        if (Translation.varIsIn id1 glob) then
          let tmpVar = 
          match tmp with
            Some(v) -> v (* this is guaranteed because we have checked it before call here. This is only to make parrent matching good *)
            | None -> failwith "Never happens!"
          in
            Comp(Labels(Pre(-1), Post(-1)), Callbool(labels, tmpVar, id2, substituteVarInExprs eList variableSubs), Assignment(Labels(Pre(-1), Post(-1)),[ (substituteVar id1 variableSubs) ], [ Variable(tmpVar) ]))
        else
          Callbool(labels, substituteVar id1 variableSubs, id2, substituteVarInExprs eList variableSubs)
    | Assignment(labels, idL, eList) -> Assignment(labels, substituteVars idL variableSubs, substituteVarInExprs eList variableSubs)
    | While(labels, d, s) -> While(labels, substituteDecider d variableSubs, prepStm variableSubs tmp glob s)
    | Returnvoid(labels) -> Returnvoid(labels)
    | Returnbool(labels, e) -> Returnbool(labels, substituteVarInExpr e variableSubs)
    | Ifelse(labels, d, s1, s2) -> Ifelse(labels, substituteDecider d variableSubs, prepStm variableSubs tmp glob s1, prepStm variableSubs tmp glob s2)
    ;;
 
(* Determine if the current procedure needs a tmp variable. This is needed if we assign a global variable the return value of a procedure. *)
let rec needsTmp glob s = 
  match s with
    | Comp(labels, s1, s2) -> (needsTmp glob s1) || (needsTmp glob s2)
    | Callbool(labels, id1, id2, eList) -> (* we are not allowed to assign a global variable the result of the procedure *)
        if (Translation.varIsIn id1 glob) then
          true
        else
          false
    | While(labels, d, s) -> needsTmp glob s
    | Ifelse(labels, d, s1, s2) -> (needsTmp glob s1) || (needsTmp glob s2)
    | _ -> false
    ;;

(* Preprocesses one procedure, and returns a new procedure *)
let prepProc proc glob = 
  let Procedure(labels, t, p, fp, d, s) = proc in
  let localVars = (Translation.getFormalParameters proc) @ (Translation.getLocals proc) in 
  let (variableSubs, newDecl) = prepDecl [] glob localVars d in
    if (needsTmp glob s) then
      let tmpVar = createNewVar glob localVars in
      let newDecl = Declcomp(Labels(Pre(-1), Post(-1)), Declaration(Labels(Pre(-1), Post(-1)), tmpVar), newDecl) in
        Procedure(labels, t, p, fp, newDecl, prepStm variableSubs (Some(tmpVar)) glob s)
    else
      let tmp = None in
        Procedure(labels, t, p, fp, newDecl, prepStm variableSubs tmp glob s)
  ;;

(* Preprocesses a list och procedures and return a new list of procedures *)
let prep procs glob = 
  let rec iterateProcs procs newProcs = 
    resetTmpCounter ();
    match procs with
		 proc :: rest -> let newProc = prepProc proc glob in iterateProcs rest (newProc :: newProcs)
     | [] -> newProcs
  in
   List.rev (iterateProcs procs []) (* performance?, change if to slow *)
;;